#include "main.h"
#include "Socket.h"

/*
 * WinsockStart
 * - Start winsock and initialize colors
 */
BOOL WinsockStart(int nSockets)
{
	// We want winsock version 1.1
	WORD wSockVersion = MAKEWORD(1, 1);
	WSADATA wsaData;

	/*
	 * WSAStartup() must be called to initialize winsock!
	 */
	int nRet = WSAStartup(wSockVersion, &wsaData);
	if (nRet != 0)
	{
		MessageBox(NULL, "Failed to start winsock.", "Error", MB_ICONERROR | MB_OK);
		return FALSE;
	}

	/*
	 * Create memory for an array of colors for our users
	 */
	sockets = (socketType *)calloc(nSockets, sizeof(socketType));
	colors = (COLORREF *)calloc(COLORCOUNT, sizeof(COLORREF));

	colors[0] = RGB(0, 0, 255);
	colors[1] = RGB(255, 0, 0);
	colors[2] = RGB(0, 255, 0);

	return TRUE;
}

/*
 * WinsockConnect
 * - Connect to another computer
 */
BOOL WinsockConnect(HWND hWnd, SOCKADDR_IN *serverInfo, char *userID)
{
	/*
	 * connect =P (kinda self explanatory)
	 */
	int nRet = connect(sockets[0].s, (LPSOCKADDR)serverInfo, sizeof(struct sockaddr));
	if (nRet == SOCKET_ERROR)
	{
		if (WSAGetLastError() == 10061)
			return FALSE;

		WinsockError(WSAGetLastError(), "WinsockConnect()");
		return FALSE;
	}

	/*
	 * Tell winsock to send us message when FD_CLOSE and FD_READ happen
	 * + FD_CLOSE - Socket is closed
	 * + FD_READ - Socket has data that is ready to be received
	 */
	WSAAsyncSelect(sockets[0].s, hWnd, WM_SOCKET_ASYNC, FD_CLOSE | FD_READ);

	/*
	 * Set winsock user ID and color.
	 * These actually arent used (for the "0" socket), but for consistency, they are here
	 */
	sockets[0].szUserID = userID;
	sockets[0].crText = colors[nCurrentColor++];

	if (nCurrentColor >= COLORCOUNT)
		nCurrentColor = 0;

	return TRUE;
}

/*
 * WinsockEnd
 * - Close all sockets and end winsock
 */
BOOL WinsockEnd(int nMaxUsers)
{
	/*
	 * Go through each socket and close it... free user too
	 */
	for (int i = 0; i < nMaxUsers; i++)
	{
		if (sockets[i].s != NULL && closesocket(sockets[i].s) == SOCKET_ERROR)
		{
			WinsockError(WSAGetLastError(), "WinsockEnd()");
			return FALSE;
		}

		if (sockets[i].szUserID != NULL)
			free(sockets[i].szUserID);

		sockets[i].s = NULL;
		sockets[i].szUserID = NULL;
	}

	free(sockets);
	free(colors);

	/*
	 * WSACleanup called to end winsock
	 */
	int nRet = WSACleanup();
	if (nRet != 0)
	{
		MessageBox(NULL, "Failed to end winsock.", "Error", MB_ICONERROR | MB_OK);
		return FALSE;
	}
	return TRUE;
}

/*
 * WinsockBroadcast
 * - Broadcast all users to ... all users. 
 */
BOOL WinsockBroadcast(int nMaxUsers)
{
	for (int i = 1; i < nMaxUsers; i++)
	{
		if (sockets[i].s == NULL)
			continue;

		/*
		 * Send CLER command to current socket.
		 * CLER tells the client to clear their connected users list
		 * The data "CLEAR" is just dummy text, it can be anything really.
		 */
		ChatSend("CLER", "CLEAR", sockets[i].s);

		for (int j = 0; j < nMaxUsers; j++)
		{
			if (sockets[j].s == NULL)
				continue;

			/*
			 * Send USER command to current socket
			 * USER tells the client to add a new user to the list
			 */
			ChatSend("USER", sockets[j].szUserID, sockets[i].s);
		}
	}

	return TRUE;
}

/*
 * WinsockOpen
 * - Creates initial socket
 */
BOOL WinsockOpen()
{
	/*
	 * Create a socket
	 */
	sockets[0].s = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (sockets[0].s == INVALID_SOCKET)
	{
		WinsockError(WSAGetLastError(), "WinsockOpen()");
		return FALSE;
	}

	return TRUE;
}

/*
 * WinsockClose
 * - Closes specific or all sockets
 */
BOOL WinsockClose(int nMaxUsers, SOCKET s)
{
	if (s == NULL)
	{
		/*
		 * Close all sockets
		 */
		for (int i = 0; i < nMaxUsers; i++)
		{
			if (sockets[i].s != NULL && closesocket(sockets[i].s) == SOCKET_ERROR)
			{
				WinsockError(WSAGetLastError(), "WinsockClose()");
				return FALSE;
			}

			if (sockets[i].szUserID != NULL)
				free(sockets[i].szUserID);

			sockets[i].s = NULL;
			sockets[i].szUserID = NULL;
		}
	}
	else
	{
		/*
		 * Close specified socket
		 */
		for (int i = 0; i < nMaxUsers; i++)
		{
			if (sockets[i].s == s)
				break;
		}

		if (i == nMaxUsers)
			return FALSE;

		if (sockets[i].s != NULL && closesocket(sockets[i].s) == SOCKET_ERROR)
		{
			WinsockError(WSAGetLastError(), "WinsockClose()");
			return FALSE;
		}

		if (sockets[i].szUserID != NULL)
			free(sockets[i].szUserID);

		sockets[i].s = NULL;
		sockets[i].szUserID = NULL;
	}

	return TRUE;
}

/*
 * WinsockGetUser
 * - Get user info for a socket
 */
BOOL WinsockGetUser(SOCKET s, int nMaxUsers, char **szUserID, COLORREF *crText)
{
	for (int i = 0; i < nMaxUsers; i++)
	{
		if (sockets[i].s == s)
			break;
	}

	if (i == nMaxUsers)
		return FALSE;

	if (szUserID != NULL && sockets[i].szUserID != NULL)
		*szUserID = sockets[i].szUserID;
	if (crText != NULL)
		*crText = sockets[i].crText;

	return TRUE;
}

/*
 * WinsockSetUser
 * - Set user ID for a socket
 */
BOOL WinsockSetUser(SOCKET s, int nMaxUsers, char *szUserID)
{
	for (int i = 0; i < nMaxUsers; i++)
	{
		if (sockets[i].s == s)
			break;
	}

	if (i == nMaxUsers)
		return FALSE;

	/*
	 * Set socket user ID and color!
	 */
	sockets[i].szUserID = szUserID;
	sockets[i].crText = colors[nCurrentColor++];

	if (nCurrentColor >= COLORCOUNT)
		nCurrentColor = 0;
	
	return TRUE;
}

/*
 * WinsockGetColor
 * - Get the color for a specific user (unknown socket)
 */
BOOL WinsockGetColor(const char *szName, int nMaxUsers, COLORREF *color)
{
	for (int i = 0; i < nMaxUsers; i++)
	{
		if (sockets[i].szUserID != NULL && strcmp(sockets[i].szUserID, szName) == 0)
		{
			/*
			 * Fill *color (value of pointer to color) with COLORREF (Packed RGB)
			 * Then return TRUE so we don't continue to reiterate through the sockets for no reason
			 */
			*color = sockets[i].crText;
			return TRUE;
		}
	}
	return FALSE;
}

/*
 * WinsockSetColor
 * - Set color for a socket
 */
BOOL WinsockSetColor(SOCKET s, int nMaxUsers, COLORREF color)
{
	for (int i = 0; i < nMaxUsers; i++)
	{
		if (sockets[i].s == s)
			break;
	}

	if (i == nMaxUsers)
		return FALSE;

	sockets[i].crText = color;
	return TRUE;
}

/*
 * WinsockSend
 * - Send data to socket
 */
BOOL WinsockSend(const char *text, SOCKET s = NULL)
{
	if (s == NULL)
		s = sockets[0].s;

	int nRet = send(s, text, strlen(text), 0);
	if (nRet == SOCKET_ERROR)
	{
		WinsockError(WSAGetLastError(), "WinsockSend()");
		return FALSE;
	}
	return TRUE;
}

/*
 * WinsockSendAll
 * - Send data to all connected sockets (except "0" socket or listening socket)
 */
BOOL WinsockSendAll(const char *szSend, int nMaxUsers)
{
	for (int i = 1; i < nMaxUsers; i++)
	{
		if (sockets[i].s == NULL)
			continue;

		int nRet = send(sockets[i].s, szSend, strlen(szSend), 0);
		if (nRet == SOCKET_ERROR)
		{
			WinsockError(WSAGetLastError(), "WinsockSendAll()");
			return FALSE;
		}
	}
	return TRUE;
}

/*
 * WinsockReadType
 * - Read first 4 bytes of incoming data (the command)
 */
BOOL WinsockReadType(SOCKET s, char *szType)
{
	/*
	 * Read first 4 bytes
	 */
	int nRet = recv(s, szType, 4, 0);
	if (nRet == SOCKET_ERROR)
	{
		if (WSAGetLastError() == 10035)
			return FALSE;

		WinsockError(WSAGetLastError(), "WinsockReadType()");
		return FALSE;
	}

	/*
	 * Add NULL character at end since recv does not (with good reason, since NULL was not part
	 * of the incoming data)
	 */
	szType[4] = '\0';
	return TRUE;
}

/*
 * WinsockReadData
 * - Read the length of the data and then read all of the data
 */
BOOL WinsockReadData(SOCKET s, char **data)
{
	/*
	 * Read the length
	 */
	char *szLength = (char *)malloc(5);
	int nRet = recv(s, szLength, 4, 0);
	if (nRet == SOCKET_ERROR)
	{
		if (WSAGetLastError() == 10035)
			return FALSE;

		WinsockError(WSAGetLastError(), "WinsockReadData()");
		return FALSE;
	}
	
	/*
	 * Add NULL character at end since recv does not (with good reason, since NULL was not part
	 * of the incoming data)
	 */
	szLength[4] = '\0';

	/*
	 * Convert length to integer
	 */
	int nLength = atoi(szLength);
	free(szLength);

	/*
	 * Receive rest of data
	 */
	char *szData = (char *)malloc(nLength + 1);
	nRet = recv(s, szData, nLength, 0);
	if (nRet == SOCKET_ERROR)
	{
		if (WSAGetLastError() == 10035)
			return FALSE;

		WinsockError(WSAGetLastError(), "WinsockReadData()");
		return FALSE;
	}

	/*
	 * Add NULL character at end since recv does not (with good reason, since NULL was not part
	 * of the incoming data)
	 */
	szData[nLength] = '\0';

	/*
	 * Set pointer to pointer to szData memory address
	 */
	*data = szData;

	return TRUE;
}

/*
 * WinsockListen
 * - Begin listening on socket "0"
 */
BOOL WinsockListen(HWND hWnd, int nPort, char *szUserID)
{
	/*
	 * Set "0" socket user data (not used, but included for consistency)
	 */
	sockets[0].szUserID = szUserID;
	sockets[0].crText = colors[nCurrentColor++];

	if (nCurrentColor >= COLORCOUNT)
		nCurrentColor = 0;

	/*
	 * Bind local address to socket
	 */
	SOCKADDR_IN serverInfo;

	serverInfo.sin_family = AF_INET;
	serverInfo.sin_addr.s_addr = INADDR_ANY;
	serverInfo.sin_port = htons(nPort);

	int nRet = bind(sockets[0].s, (LPSOCKADDR)&serverInfo, sizeof(struct sockaddr));
	if (nRet == SOCKET_ERROR)
	{
		WinsockError(WSAGetLastError(), "WinsockListen()");
		return FALSE;
	}

	/*
	 * Start listening
	 */
	nRet = listen(sockets[0].s, 5);
	if (nRet == SOCKET_ERROR)
	{
		WinsockError(WSAGetLastError(), "WinsockListen()");
		return FALSE;
	}

	/*
	 * Tell winsock to send us messages when FD_ACCEPT, FD_READ, AND FD_CLOSE happen
	 * + FD_ACCEPT - Connection is accepted
	 * + FD_READ - Data is ready to be received
	 * + FD_CLOSE - Socket was closed
	 */
	WSAAsyncSelect(sockets[0].s, hWnd, WM_SOCKET_ASYNC, FD_ACCEPT | FD_READ | FD_CLOSE);

	return TRUE;
}

/*
 * WinsockAccept
 * - Accept connection request
 */
BOOL WinsockAccept(int nMaxUsers)
{
	/*
	 * Find open socket
	 */
	for (int i = 0; i < nMaxUsers; i++)
	{
		if (sockets[i].s == NULL)
			break;
	}

	if (i == nMaxUsers)
		return FALSE;

	/*
	 * Accept connection
	 */
	sockets[i].s = accept(sockets[0].s, NULL, NULL);
	if (sockets[i].s == SOCKET_ERROR)
	{
		WinsockError(WSAGetLastError(), "WinsockAccept()");
		return FALSE;
	}

	return TRUE;
}

/*
 * WinsockError
 * - Raise error
 */
void WinsockError(int nSocketError, const char *szFunction)
{
	char cErrorMsg[92];

	ZeroMemory(cErrorMsg, 92);
	sprintf(cErrorMsg, "Call to %s returned error %d!", szFunction, nSocketError);

	MessageBox(NULL, cErrorMsg, "Socket Error", MB_ICONERROR | MB_OK);
}
